from neo4j import GraphDatabase


class Neo4jConnection:
    def __init__(self):
        self.uri = "neo4j://localhost:7687"
        self.username = "neo4j"
        self.password = "JGL12345"
        self.driver = GraphDatabase.driver(self.uri, auth=(self.username, self.password))

    def execute_query(self, query, params):
        with self.driver.session() as session:
            result = session.run(query, params)
            return [record for record in result]

    def query_persons(self):
        query = "match (p:Person) return p.name as name"
        return self.execute_query(query, {})

    def query_places(self):
        query = "match (p:Place) return p.name as name"
        return self.execute_query(query, {})

    def query_houses(self):
        query = "match (h:House) return h.name as name"
        return self.execute_query(query, {})

    def query_events(self):
        query = "match (e:Event) return e.name as name"
        return self.execute_query(query, {})


graph = Neo4jConnection()
entity_type = "person"

relation_types = [
    "朋友", "友人", "好友", "密友", "挚友",
    "敌人", "对手", "大敌", "敌对方",
    "大臣", "臣子", "手下的官员",
    "手下", "部下", "下属"
]

sources = [
    "[*](^)有哪些[$](#)？",
    "[*](^)的[$](#)都有谁？",
    "[*](^)的[$](#)都是些什么人？",
    "你知道[*](^)的[$](#)都有谁吗？",
    "你了解[*](^)的[$](#)吗？",
    "[*](^)拥有哪些[$](#)？"
]

examples = []

# with open("other_intents.txt", "r", encoding="utf-8") as file:
#    for line in file:
#        examples.append(line)

records = graph.query_persons()

for record in records:
    name = record["name"]
    for source in sources:
        person_replaced = source.replace("*", name).replace("^", entity_type)
        for relation in relation_types:
            tmp = person_replaced.replace("$", relation).replace("#", "relation")
            examples.append(tmp)

with open("tmp.in", "w", encoding="utf-8") as file:
    for example in examples:
        file.write("    - " + example + '\n')
